"""
均线模块
负责计算和绘制移动平均线
"""

import pandas as pd
import numpy as np
import pyqtgraph as pg
from typing import List, Tuple, Dict, Optional
from PyQt5 import QtCore, QtGui
from .BaseChart import BaseChart
from .config import config, flags, logger


class MovingAverageCalculator:
    """移动平均线计算器"""
    
    def __init__(self, data: pd.DataFrame):
        """
        初始化移动平均线计算器
        
        Args:
            data: 包含股票数据的DataFrame
        """
        self.data = data.copy()
        self.ma_data = {}  # 存储计算好的均线数据
    
    def calculate_ma(self, periods: List[int], price_column: str = 'close') -> Dict[int, pd.Series]:
        """
        计算指定周期的移动平均线
        
        Args:
            periods: 均线周期列表，如 [5, 10, 20, 30]
            price_column: 价格列名，默认为'close'
            
        Returns:
            包含各周期均线的字典
        """
        ma_results = {}
        for period in periods:
            if period <= 0:
                continue
                
            # 计算移动平均线
            ma_series = self.data[price_column].rolling(window=period, min_periods=1).mean()
            ma_results[period] = ma_series
            
            # 存储到实例变量中
            self.ma_data[period] = ma_series
        
        return ma_results
    
    def get_ma_data(self, period: int) -> Optional[pd.Series]:
        """获取指定周期的均线数据"""
        return self.ma_data.get(period)
    
    def get_ma_range(self, periods: List[int]) -> Tuple[float, float]:
        """
        获取均线的价格范围
        
        Args:
            periods: 均线周期列表
            
        Returns:
            (最小价格, 最大价格)
        """
        min_price = float('inf')
        max_price = float('-inf')
        
        for period in periods:
            if period in self.ma_data:
                ma_series = self.ma_data[period]
                min_price = min(min_price, ma_series.min())
                max_price = max(max_price, ma_series.max())
        
        return min_price, max_price


class MovingAverageItem(pg.GraphicsObject):
    """移动平均线绘制项"""
    
    def __init__(self, ma_data: pd.Series, color: str = 'white', width: int = 1):
        """
        初始化移动平均线绘制项
        
        Args:
            ma_data: 均线数据
            color: 线条颜色
            width: 线条宽度
        """
        pg.GraphicsObject.__init__(self)
        self.ma_data = ma_data
        self.color = color
        self.width = width
        self.generatePicture()
    
    def generatePicture(self):
        """生成均线图形"""
        self.picture = QtGui.QPicture()
        p = QtGui.QPainter(self.picture)
        
        # 设置线条样式 - 使用更安全的颜色设置
        try:
            pen = pg.mkPen(self.color, width=self.width)
        except:
            # 如果颜色解析失败，使用默认颜色
            pen = pg.mkPen('white', width=self.width)
        p.setPen(pen)
        
        # 绘制均线
        points = []
        for i, value in enumerate(self.ma_data):
            if not pd.isna(value):
                points.append(QtCore.QPointF(i, value))
        
        if len(points) > 1:
            # 绘制连续线条
            for i in range(len(points) - 1):
                p.drawLine(points[i], points[i + 1])
        
        p.end()
    
    def paint(self, painter, option, widget):
        """绘制均线"""
        painter.drawPicture(0, 0, self.picture)
    
    def boundingRect(self):
        """返回边界矩形"""
        return QtCore.QRectF(self.picture.boundingRect())


class MovingAverageManager(BaseChart):
    """移动平均线管理器"""
    
    def __init__(self, plot_widget: pg.PlotWidget):
        """
        初始化移动平均线管理器
        
        Args:
            plot_widget: PyQtGraph绘图控件
        """
        if flags.ENABLE_MOVING_AVERAGE:
            # self.plot_widget = plot_widget
            super().__init__(plot_widget)
            self.ma_items: Dict[int, MovingAverageItem] = {}
            self.ma_calculator: Optional[MovingAverageCalculator] = None
            
            # 默认均线配置
            self.default_periods = [5, 10, 20, 30]
            self.default_colors = {
                5: 'y',      # 5日均线 - 黄色
                10: 'm',     # 10日均线 - 紫色
                20: 'b',     # 20日均线 - 蓝色
                30: 'r'      # 30日均线 - 红色
            }
    
    def set_data(self, data: pd.DataFrame):
        """
        设置数据并初始化计算器
        
        Args:
            data: 股票数据DataFrame
        """
        if not flags.ENABLE_MOVING_AVERAGE:
            return

        self.ma_calculator = MovingAverageCalculator(data)
    
    def add_moving_averages(self, periods: List[int] = None, 
                          colors: Dict[int, str] = None,
                          price_column: str = 'close'):
        """
        添加移动平均线
        
        Args:
            periods: 均线周期列表，默认为[5, 10, 20, 30]
            colors: 各周期对应的颜色字典
            price_column: 价格列名
        """
        if not flags.ENABLE_MOVING_AVERAGE:
            return
        if self.ma_calculator is None:
            raise ValueError("请先调用 set_data() 方法设置数据")
        
        # 使用默认配置
        if periods is None:
            periods = self.default_periods
        if colors is None:
            colors = self.default_colors
        
        # 计算均线
        ma_results = self.ma_calculator.calculate_ma(periods, price_column)
        
        # 移除旧的均线
        self.remove_all_moving_averages()
        
        # 添加新的均线
        for period in periods:
            if period in ma_results:
                print(f"添加均线: {period}")
                ma_data = ma_results[period]
                color = colors.get(period, 'white')
                
                # 创建均线绘制项
                ma_item = MovingAverageItem(ma_data, color)
                self.ma_items[period] = ma_item
                
                # 添加到图表
                self.plot_widget.addItem(ma_item)
                
                # 设置层级（在K线下方）
                ma_item.setZValue(50)
    
    def remove_moving_average(self, period: int):
        """
        移除指定周期的移动平均线
        
        Args:
            period: 均线周期
        """
        if not flags.ENABLE_MOVING_AVERAGE:
            return
        if period in self.ma_items:
            self.plot_widget.removeItem(self.ma_items[period])
            del self.ma_items[period]
    
    def remove_all_moving_averages(self):
        """移除所有移动平均线"""
        if not flags.ENABLE_MOVING_AVERAGE:
            return

        for period in list(self.ma_items.keys()):
            self.remove_moving_average(period)
    
    def update_moving_averages(self, periods: List[int] = None,
                              colors: Dict[int, str] = None,
                              price_column: str = 'close'):
        """
        更新移动平均线
        
        Args:
            periods: 均线周期列表
            colors: 各周期对应的颜色字典
            price_column: 价格列名
        """
        if not flags.ENABLE_MOVING_AVERAGE:
            return
        self.add_moving_averages(periods, colors, price_column)
    
    def get_ma_range(self, periods: List[int] = None) -> Tuple[float, float]:
        """
        获取均线的价格范围
        
        Args:
            periods: 均线周期列表
            
        Returns:
            (最小价格, 最大价格)
        """
        if self.ma_calculator is None:
            return 0.0, 0.0
        
        if periods is None:
            periods = list(self.ma_items.keys())
        
        return self.ma_calculator.get_ma_range(periods)
    
    def get_ma_data(self, period: int) -> Optional[pd.Series]:
        """获取指定周期的均线数据"""
        if self.ma_calculator is None:
            return None
        return self.ma_calculator.get_ma_data(period)
    
    def get_available_periods(self) -> List[int]:
        """获取可用的均线周期列表"""
        return list(self.ma_items.keys())
    
    def set_ma_visible(self, period: int, visible: bool):
        """
        设置指定周期均线的可见性
        
        Args:
            period: 均线周期
            visible: 是否可见
        """
        if not flags.ENABLE_MOVING_AVERAGE:
            return
        if period in self.ma_items:
            self.ma_items[period].setVisible(visible)
    
    def set_all_ma_visible(self, visible: bool):
        """
        设置所有均线的可见性
        
        Args:
            visible: 是否可见
        """
        if not flags.ENABLE_MOVING_AVERAGE:
            return
        for ma_item in self.ma_items.values():
            ma_item.setVisible(visible)

    def set_x_range(self, x_min: float, x_max: float):
        """
        设置X轴范围
        
        Args:
            x_min: X轴最小值
            x_max: X轴最大值
        """
        if not flags.ENABLE_MOVING_AVERAGE:
            return
        super().set_x_range(x_min, x_max)
    
    def set_date_ticks(self, date_ticks: List[Tuple[int, str]]):
        """
        设置X轴日期刻度
        
        Args:
            date_ticks: 日期刻度列表
        """
        if not flags.ENABLE_MOVING_AVERAGE:
            return
        super().set_date_ticks(date_ticks)

    def set_price_range(self, min_price: float, max_price: float, padding: float = 0.1):
        """
        设置价格范围
        
        Args:
            min_price: 最低价格
            max_price: 最高价格
            padding: 边距比例
        """
        if not flags.ENABLE_MOVING_AVERAGE:
            return
        super().set_price_range(min_price, max_price, padding)